import ismechanisms.ismech
import ismechanisms.istools
import ismechanisms.isexp
import argparse
from mbi import Dataset, Domain
import pandas as pd
import numpy as np
import itertools
import os
import json
import IncreSyn

config_path = "./config.json"
if os.path.exists(config_path):
    config_args = IncreSyn.load_config(config_path=config_path)
    args = config_args
    print("Loaded config from config file!") 
else:
    print("Config file not found! Exit.")
    exit()

# Universal part start
data = Dataset.load(args.dataset, args.domain)
# IncreSyn: Pre-parse mech
# args.mech = args.mech.lower
prng = np.random
# IncreSyn: Prepare Workload
workload = list(itertools.combinations(data.domain, args.degree))
workload = [cl for cl in workload if data.domain.size(cl) <= args.max_cells]
if args.num_marginals is not None:
    workload = [workload[i] for i in prng.choice(len(workload), args.num_marginals, replace=False)]
if args.mech == "aim":
    workload = [(cl, 1.0) for cl in workload]
    
# IncreSyn: Load last synthetic data
if args.lastsyn is not None:
    lastsyn_load = Dataset.load(args.lastsyn, args.domain)
else:
    lastsyn_load = None
# Universal part end

attr_name = []
attr_name.append("Dataset")
attr_name.append("Budgets")
attr_name.append("UniversalError")
attr_name.append("PreferedError")
attr_name.append("TimeConsume")
log_file_name = ismechanisms.istools.log_init(attr_name=attr_name)
log_file = ismechanisms.istools.info_logger("======Experiment Type A START======")
dataset_name = args.name

budget_list = [0.5, 0.75, 1.0, 1.25, 1.5]
i=1
for budget in budget_list:
    mech_para = ismechanisms.isexp.args_handler(args, budgets=budget,log_file=log_file) # 这里的epsilon是为了方便后续实验使用不同的预算输入
    exp_results = [] 
    print("Preparing original data round "+str(i)+"/"+str(len(budget_list))+"...",end="",flush=True)
    exp_results = ismechanisms.isexp.original_syn(mech_para=mech_para, mech_type=args.mech, data=data,workload=workload,error_method=args.error_method)
    exp_results.insert(0, str(budget))
    exp_results.insert(0, dataset_name+"Original")
    ismechanisms.istools.log_append(exp_results, log_file_name[0], log_file_name[1])
    print("Done")

    print("Starting exp round round "+str(i)+"/"+str(len(budget_list))+"...",end="",flush=True)
    exp_results = ismechanisms.isexp.exp_single(mech_para=mech_para, mech_type=args.mech, data=data,workload=workload,error_method=args.error_method)
    exp_results.insert(0, str(budget))
    exp_results.insert(0, dataset_name+"Optimized")
    ismechanisms.istools.log_append(exp_results, log_file_name[0], log_file_name[1])
    print("Done")
    i+=1
    